import torch
import numpy as np
import cv2
import torch





def calaccurary(tlabel,label):
    ''' 计算不同类别的准确率'''
    tlabel=np.argmax(tlabel.cpu().numpy(),axis=1)
    label=label.cpu().numpy()
    if len(label)==1:
        if label[0]==tlabel[0]:
            return [1,1,1]
        else:
            return [0,0,0]
    res=label-tlabel # 计算多少类别不对
    
    cs=len(np.argwhere(res==0))/len(label) # 表示类别全对
    a0=-1
    if len(np.argwhere(label==0))==0:
        a0=-1
    else:
        a0=1-len(np.argwhere(res==-1))/len(np.argwhere(label==0))
    a1=-1
    if len(np.argwhere(label==1))==0:
        a1=-1
    else:
        a1=1-len(np.argwhere(res==1))/len(np.argwhere(label==1))
    return [cs,a0,a1]


def visualize_cam(mask, img):
    """Make heatmap from mask and synthesize GradCAM result image using heatmap and img.
    Args:
        mask (torch.tensor): mask shape of (1, 1, H, W) and each element has value in range [0, 1]
        img (torch.tensor): img shape of (1, 3, H, W) and each pixel value is in range [0, 1]
        
    Return:
        heatmap (torch.tensor): heatmap img shape of (3, H, W)
        result (torch.tensor): synthesized GradCAM result of same shape with heatmap.
    """
    heatmap = cv2.applyColorMap(np.uint8(255 * mask.squeeze()), cv2.COLORMAP_JET)
    heatmap = torch.from_numpy(heatmap).permute(2, 0, 1).float().div(255)
    b, g, r = heatmap.split(1)
    heatmap = torch.cat([r, g, b])
    
    result = heatmap+img.cpu()
    result = result.div(result.max()).squeeze()
    
    return heatmap, result

def outcalInfoOfConv(calmode):
    pass